import os
import re
import signal
from tempfile import NamedTemporaryFile, mkstemp
from threading import Thread, Event, Lock
from subprocess import Popen, PIPE
# TODO: look into the pty module and openpty(), to address some problems we've seen with popen and subprocess
from getpass import getpass
import paramiko
from async import AsyncIterable, AsyncIterableGroup, AsyncThreadException
from support import PickledExecution, ASYNCTRACE

def _bytes_from_stream(stream):
    """yields the bytes of a stream one at a time"""
    next_char = stream.read(1)
    while next_char:
        yield next_char
        next_char = stream.read(1)

def create_connection(username, hostname, port = 22, getpass = None, su_source_username = None, su_source_getpass = None):

        if hostname == "localhost":
            # we have a local scenario
        
            if username == os.getlogin():
                # local user, on local machine
                return LocalUserConnection()
            
            else:
                # other user, on local machine
                raise UnsupportedConnectionError("currently cannot support local users other than the current user (use os.getlogin() to get the local username)")
    
        else:
            # we have a remote scenario
            if su_source_username:
                # user on remote machine who needs to be SSHed in
                # and then su'd to before it can do anything
                su_from_connection = RemoteConnection(su_source_username, hostname, getpass = su_source_getpass, port = port)
                return RemoteSuConnection(
                    su_from_connection = su_from_connection,
                    username = username,
                    getpass = getpass,
                )
        
            else:
                # this is a user who can remotely log in to the machine
                # and execute commmands
                return RemoteConnection(
                    username = username,
                    hostname = hostname,
                    getpass = getpass,
                    port = port,
                )


class SFTPFileAdapter(object):
    
    def __init__(self, sftp_file, filepath):
        self.__sftp_file = sftp_file
        self.__filepath = filepath
    
    def __getattr__(self, attr_name):
        if attr_name == "name":
            return self.__filepath
        return getattr(self.__sftp_file, attr_name)

class GeneratorStream(object):
    
    def __init__(self, generator, input_handler, kill_handler):
        self.__async_iterable = AsyncIterable( generator )
        self.__input_handler = input_handler
        self.__kill_handler = kill_handler
        self.__input_lock = Lock()
    
    def __iter__(self):
        for obj in self.__async_iterable:
            yield obj
    
    def inject(self, obj):
        # TODO: this may need some kind of special object (None may not be appropriate) to indicate when no more input will be provided?
        self.__input_lock.acquire()
        try:
            self.__input_handler( obj )
        finally:
            self.__input_lock.release()
        return self
    
    def drain(self):
        for dontcare in self:
            pass
    
    def kill(self):
        self.__kill_handler()


class CommandStream(object):
    
    def __init__(self, stdout_generator, stderr_generator, input_handler, kill_handler):
        self.__async_iterable = AsyncIterableGroup( stdout_generator, stderr_generator )
        self.__input_handler = input_handler
        self.__kill_handler = kill_handler
        self.__input_lock = Lock()

    def __iter__(self):
        """yields each byte of the stream, and also
        (as a convenience) each line of the stream"""
        pending_stdout_line = ""
        pending_stderr_line = ""
        for stdout_byte, stderr_byte in self.__async_iterable:
            stdout_line = None
            stderr_line = None
            if stdout_byte:
                if stdout_byte == "\n":
                    stdout_line = pending_stdout_line
                    pending_stdout_line = ""
                else:
                    pending_stdout_line += stdout_byte
            if stderr_byte:
                if stderr_byte == "\n":
                    stderr_line = pending_stderr_line
                    pending_stderr_line = ""
                else:
                    pending_stderr_line += stderr_byte
            yield stdout_line, stderr_line, stdout_byte, stderr_byte
    
    def inject(self, buffer_str):
        self.__input_lock.acquire()
        try:
            self.__input_handler( buffer_str )
        finally:
            self.__input_lock.release()
        return self
    
    def drain(self):
        for dontcare in self:
            pass
    
    def kill(self):
        self.__kill_handler()


class UnsupportedConnectionError(Exception): pass
class AuthenticationError(Exception): pass


class NonzeroReturnStatus(Exception):
    
    code = property(lambda self: self.__code)
    
    def __init__(self, code, command):
        Exception.__init__(self, "'%s' returned exit status %s" % (command, code))
        self.__code = code


class AbstractConnection(object):
    
    def create_command_stream(self, command):
        """return an iterable that sequences over (stdout, stderr)"""
        raise NotImplementedError
        
    def create_generator_stream(self, generator, virtualenv):
        """return an async iterable that sequences over (obj) and raises
        any remote exceptions as PickledExecutionExceptions"""
        raise NotImplementedError
    
    def open_file(self, filepath, mode = "w+"):
        """opens a stream to the given file on the connection"""
        raise NotImplementedError
    
    def open_tempfile(self):
        """opens a temp file on the given connection"""
        raise NotImplementedError


class LocalUserConnection(AbstractConnection):
    
    def __init__(self):
        AbstractConnection.__init__(self)
    
    def open_file(self, filepath, mode = "w+"):
        return open(filepath, mode)
    
    def open_tempfile(self):
        filename = mkstemp()[1]
        return self.open_file(filename)
    
    def create_command_stream(self, command):

        pipe = Popen([command], shell = True, close_fds = True, stdout = PIPE, stderr = PIPE, stdin = PIPE)
        done_event = Event()
        
        def watch_stdout():
            
            # yield each line from the stream
            for byte in _bytes_from_stream(pipe.stdout):
                yield byte
                
            # raise exceptions on bad return codes
            done_event.wait()
            code = pipe.returncode
            if code != 0:
                raise NonzeroReturnStatus(code, command)

        def watch_stderr():
            
            # yield each line from the stream
            for byte in _bytes_from_stream(pipe.stderr):
                yield byte
                
            # raise exceptions on bad return codes
            done_event.wait()
            code = pipe.returncode
            if code != 0:
                raise NonzeroReturnStatus(code, command)
        
        # hack to reliably watch for the process to finish
        # (for some reason, pipe.wait() could sometimes fail with no child process)
        def watch_status():
            try:
                if pipe.returncode is None:
                    pipe.wait()
            except OSError, e:
                # TODO: should this do anything? the process is done at this point either way
                pass
            done_event.set()
        
        Thread(target = watch_status).start()
        
        # set up an input handler
        def input_handler(buffer_str):
            pipe.stdin.write( buffer_str )
            
        def kill_handler():
            os.kill( pipe.pid, signal.SIGKILL )

        # FIXME: somehow needs to register an 'abort' method that cuts off all streams
        return CommandStream(watch_stdout, watch_stderr, input_handler, kill_handler)
        
    def create_generator_stream(self, generator, virtualenv = None):
        
        # create the pickled execution for this generator
        pickled_execution = PickledExecution(generator)
        
        # write the code that must be executed
        temp_code_file = NamedTemporaryFile()
        temp_code_file.write( pickled_execution.code )
        temp_code_file.flush()
        
        # execute it via the shell
        if virtualenv == None:
            python_cmd = "python %s" % temp_code_file.name
        else:
            python_cmd = "%s/bin/python %s" % (virtualenv, temp_code_file.name)
        pipe = Popen([python_cmd], shell = True, close_fds = True, stdout = PIPE, stderr = PIPE, stdin = PIPE)
        
        # asynchronously make sure the temp file is deleted after finishing the command
        def close_tempfile_on_completion():
            pipe.wait()
            temp_code_file.close()
            
        Thread(target = close_tempfile_on_completion).start()
        
        # set up an input handler
        def input_handler(obj):
            pickled_execution.write_object_to_stdin(obj, pipe.stdin)
            
        def kill_handler():
            os.kill( pipe.pid, signal.SIGKILL )
        
        # do the iteration over the resulting python execution
        return GeneratorStream( lambda: pickled_execution.get_response_iterator(pipe.stdout, pipe.stderr), input_handler, kill_handler )


class RemoteConnection(AbstractConnection):
    
    def __init__(self, username, hostname, getpass, port):
        AbstractConnection.__init__(self)
        
        # keep track of our connection metadata
        self.__username = username
        self.__hostname = hostname
        self.__port = int(port)
        self.__getpass_callback = getpass
        
        # create our internal ssh client
        self.__ssh_client = paramiko.SSHClient()
        self.__ssh_client.load_system_host_keys()
        self.__ssh_client.set_missing_host_key_policy(paramiko.WarningPolicy)
        self.__sftp_client = None
    
    def __del__(self):
        self.__ssh_client.close()
        if self.__sftp_client:
            self.__sftp_client.close()
    
    def _ensure_sftp_is_ready(self):
        if not self.__sftp_client:
            self._ensure_connected_and_authenticated()
            try:
                self.__sftp_client = self.__ssh_client.open_sftp()
            except Exception, e:
                self.__sftp_client.close()
                raise
    
    def _ensure_connected_and_authenticated(self):
        
        # keep track of the last exception we got from paramiko during authentication attempts
        last_ssh_exception = None
        
        # only do this if we are not authenticated yet
        def is_connected_and_authenticated():
            return self.__ssh_client.get_transport() and self.__ssh_client.get_transport().is_authenticated()
        
        # ATTEMPT 1: use local PKeys
        if not is_connected_and_authenticated():
            
            try:
                self.__ssh_client.connect(
                    hostname = self.__hostname,
                    username = self.__username,
                    port = self.__port,
                )
            except paramiko.SSHException, e:
                last_ssh_exception = e
                self.__ssh_client.close()
            
        # ATTEMPT 2: use ssh agent on the local machine
        if not is_connected_and_authenticated():
        
            try:
                agent = paramiko.Agent()
                
                for pkey in agent.get_keys():
                    try:
                        self.__ssh_client.connect(
                            hostname = self.__hostname,
                            username = self.__username,
                            port = self.__port,
                            pkey = pkey,
                            allow_agent = True,
                            look_for_keys = True
                        )
                    except paramiko.SSHException, e:
                        last_ssh_exception = e
                        self.__ssh_client.close()
                        
            except paramiko.SSHException, e:
                last_ssh_exception = e
                self.__ssh_client.close()
            
        # ATTEMPT 3: final attempt, use password
        if not is_connected_and_authenticated():
            
            try:
                if self.__getpass_callback:
                    password = self.__getpass_callback()
                    self.__ssh_client.connect(
                        hostname = self.__hostname,
                        username = self.__username,
                        port = self.__port,
                        password = password
                    )
            except paramiko.SSHException, e:
                last_ssh_exception = e
                self.__ssh_client.close()
        
        # final check
        if not is_connected_and_authenticated() and last_ssh_exception:
            # re-raise if we *still* aren't authenticated
            raise AuthenticationError("failed to authenticate for '%(username)s@%(hostname)s' after trying accessible ssh keys, sshagent, and password callback. (reason: %(paramiko_reason)s)" % dict(
                username = self.__username,
                hostname = self.__hostname,
                paramiko_reason = last_ssh_exception.message,
            ))
    
    def open_file(self, filepath, mode = "w+"):
        self._ensure_sftp_is_ready()
        sftp_file = self.__sftp_client.file( filepath, mode )
        return SFTPFileAdapter(sftp_file, filepath)
    
    def open_tempfile(self):
        temp_filename = None
        for stdout_line, stderr_line, stdout_byte, stderr_byte in self.create_command_stream("python -c 'import tempfile; print tempfile.mkstemp()[1]'"):
            if stdout_line:
                temp_filename = stdout_line
        
        sftp_file = self.open_file(temp_filename, "w+")
        return SFTPFileAdapter(sftp_file, temp_filename)
    
    def create_command_stream(self, command):

        self._ensure_connected_and_authenticated()
        
        # open an SSH shell
        channel = self.__ssh_client.get_transport().open_session()
        
        # get a pseudo-terminal so that certain commands (e.g. scp, rsync, su) behave more nicely
        channel.get_pty()
        
        # grab the streams on that channel
        bufsize = -1
        stdin = channel.makefile('wb', bufsize)
        stdout = channel.makefile('rb', bufsize)
        stderr = channel.makefile_stderr('rb', bufsize)

        # create a callback that will execute the command and return its status
        channel.exec_command(command)
        
        stdout_done_event = Event()
        stderr_done_event = Event()
        
        def watch_stdout():
            
            try:
                # yield each line from the stream
                for byte in _bytes_from_stream(stdout):
                    yield byte
                
                # raise exceptions on bad return codes
                ASYNCTRACE("stdout: receiving exit status")
                code = channel.recv_exit_status()
                if code != 0:
                    ASYNCTRACE("stdout: nonzero return code found, raising NonzeroReturnStatus")
                    raise NonzeroReturnStatus(code, command)
                ASYNCTRACE("stdout: no problems (status = 0)")

            finally:
            
                stdout_done_event.set()

        def watch_stderr():
            
            try:
                # yield each line from the stream
                for byte in _bytes_from_stream(stderr):
                    yield byte
                
                # raise exceptions on bad return codes
                ASYNCTRACE("stderr: receiving exit status")
                code = channel.recv_exit_status()
                if code != 0:
                    ASYNCTRACE("stderr: nonzero return code found, raising NonzeroReturnStatus")
                    raise NonzeroReturnStatus(code, command)
                ASYNCTRACE("stderr: no problems (status = 0)")
            
            finally:
                stderr_done_event.set()
                
        def close_channel_when_stdout_and_stderr_are_done():
            ASYNCTRACE("waiting to close channel")
            stdout_done_event.wait()
            stderr_done_event.wait()
            ASYNCTRACE("closing channel")
            channel.close()
            ASYNCTRACE("closed channel")
        
        # be sure to close down the channel when done
        Thread(name = "Thread-close_channel_when_stdout_and_stderr_are_done", target = close_channel_when_stdout_and_stderr_are_done).start()
        
        # set up an input handler
        def input_handler(buffer_str):
            stdin.write( buffer_str )
            
        def kill_handler():
            channel.close()
                
        # FIXME: should also put in stdin, once the IOManager code is in place.  This way you can track what input has gone to the process
        return CommandStream(watch_stdout, watch_stderr, input_handler, kill_handler)

        
    def create_generator_stream(self, generator, virtualenv = None):
        if virtualenv == None:
            python_cmd_format = "python %s"
        else:
            python_cmd_format = "%s/bin/python %s" % virtualenv
        return self._create_generator_stream_with_python_cmd(generator, python_cmd_format)
    
    def _create_generator_stream_with_python_cmd(self, generator, python_cmd_format):
        self._ensure_connected_and_authenticated()
        
        # create the pickled execution for this generator
        pickled_execution = PickledExecution(generator)
        
        # create a temporary file to write to
        temp_code_file = self.open_tempfile()
        temp_code_file.write( pickled_execution.code )
        temp_code_file.close()
        
        # execute it via the shell
        python_cmd = python_cmd_format % temp_code_file.name
        channel = self.__ssh_client.get_transport().open_session()
        channel.get_pty()
        bufsize = -1
        stdin = channel.makefile('wb', bufsize)
        stdout = channel.makefile('rb', bufsize)
        stderr = channel.makefile_stderr('rb', bufsize)
        channel.exec_command(python_cmd)
        
        # set up an input handler
        def input_handler(obj):
            pickled_execution.write_object_to_stdin(obj, stdin)
            
        def kill_handler():
            channel.close()
        
        # do the iteration over the resulting python execution
        return GeneratorStream( lambda: pickled_execution.get_response_iterator(stdout, stderr), input_handler, kill_handler )


class RemoteSuConnection(AbstractConnection):
    """a very hackish AbstractConnection that allows a user to be su'ed to"""
    
    def __init__(self, su_from_connection, username, getpass = None):
        self.__su_from_connection = su_from_connection
        self.__username = username
        self.__getpass = getpass
    
    def open_file(self, filepath, mode = "w+"):
        raise NotImplementedError("'su' functionality is experimental and cannot open files as the current user yet")
    
    def open_tempfile(self, filepath, mode = "w+"):
        raise NotImplementedError("'su' functionality is experimental and cannot open files as the current user yet")
    
    def create_command_stream(self, command):
        if "\"" in command:
            raise ValueError("RemoteSuConnection does not yet support the command '%s'" % command)
        
        # wrap the command and execute it on the su_from RemoteConnection
        su_wrapped_command = "su %s -c \"%s\"" % (self.__username, command)
        command_stream = self.__su_from_connection.create_command_stream(su_wrapped_command)
        
        # wrap the CommandStream we got, by watching stdout/stderr
        # for the first password prompt and responding
        password_prompt_pattern = re.compile(".*assword:.*")
        def password_enabled_stream_iteration():
            pending_stdout_line = ""
            pending_stderr_line = ""
            did_respond_to_password = False
            for stdout_line, stderr_line, stdout_byte, stderr_byte in command_stream:
                
                # track each line as it is built
                if stdout_byte:
                    if stdout_byte == "\n":
                        pending_stdout_line = ""
                    else:
                        pending_stdout_line += stdout_byte
                if stderr_byte:
                    if stderr_byte == "\n":
                        pending_stderr_line = ""
                    else:
                        pending_stderr_line += stderr_byte
                
                # respond to passwords
                if not did_respond_to_password:
                    if password_prompt_pattern.search(pending_stdout_line):
                        password = self.__getpass()
                        command_stream.inject(password + "\r\n")
                        did_respond_to_password = True
                    elif password_prompt_pattern.search(pending_stderr_line):
                        password = self.__getpass()
                        command_stream.inject(password + "\r\n")
                        did_respond_to_password = True
                
                # yield the exact set of data that we *would* have yielded
                yield stdout_line, stderr_line, stdout_byte, stderr_byte
        
        # declare a helper class that will wrap this stream
        class _SuCommandStream(object):
            def __init__(self, command_stream, password_stream_iterator):
                self.__command_stream = command_stream
                self.__password_stream_iterator = password_stream_iterator
            def __iter__(self):
                for stdout_line, stderr_line, stdout_byte, stderr_byte in self.__password_stream_iterator():
                    yield stdout_line, stderr_line, stdout_byte, stderr_byte
            def inject(self, *args, **kwargs):
                self.__command_stream.inject(*args,**kwargs)
            def drain(self, *args, **kwargs):
                self.__command_stream.drain(*args,**kwargs)
            def kill(self, *args, **kwargs):
                self.__command_stream.kill(*args,**kwargs)
                
        # return our wrapped CommandStream
        return _SuCommandStream( command_stream, password_enabled_stream_iteration )
    
    def create_generator_stream(self, generator, virtualenv = None):
        raise NotImplementedError("'su' functionality is experimental and cannot execute remote Python code yet")
    